import numpy as np
import random

class Agent():
    '''
    days: int,表示一周有几天
    num_nights: int,表示一周去几天
    utility_vectory：np.array,表示每个agent选择某一天去的时候的回报值
    last_action: np.array,每个agent上周做的选择
    '''
    def __init__(self, days, num_nights, actions):
        self.days = days
        self.num_nights = num_nights
        self.utility_vectory = np.ones(days) * 100
        self.take_random_action(actions)

    def take_random_action(self, actions):
        self.last_action = np.array(random.choice(actions))

    # 更新每个utility-vectory
    def learn(self, days, num_nights, discount_rate, reward_function, actions, attendance):
        self.utility_vectory = self.utility_vectory * discount_rate
        index = actions.index(list(self.last_action))
        reward = self.get_reward(days, num_nights, reward_function, attendance)
        self.utility_vectory[index] += reward

    # 计算每个agent的reward
    def get_reward(self, days, num_nights, reward_function, attendance):
        if reward_function == 1:
            # wonderful-life-0
            return 1 + global_utility(attendance, num_nights) \
                   - global_utility(attendance - self.last_action, num_nights)
        elif reward_function == 2:
            # wonderful-life-1
            t = np.ones(days)
            return 1 + global_utility(attendance, num_nights) \
                   - global_utility(attendance - self.last_action + t, num_nights)
        elif reward_function == 3:
            # aristocratic
            t = np.ones(days) * (1/7)
            return 1 + global_utility(attendance, num_nights) - \
                   global_utility(attendance - self.last_action + t, num_nights)
        else:
            print("input reward-function error")

    # 做出选择（哪一天去酒吧）
    def take_action(self, actions):
        self.last_action = actions[stochastic_choice(self.utility_vectory)]
        # self.last_action = actions[boltzmann_choice(self.utility_vectory)]
        # self.last_action = actions[self.utility_vectory.argmax()]


# 计算整体回报值
def global_utility(att, num_nights):
    c = [3, 6, 8, 10, 12, 15]
    return sum(att * np.exp(-att/c[num_nights-1]))


# 随机选择一天去酒吧
def stochastic_choice(vec):
    r = random.uniform(0, sum(vec))
    base_prob = 0
    for i in range(len(vec)):
        base_prob += vec[i]
        if r <= base_prob:
            return i


def boltzmann_choice(vec):
    v = np.exp(vec / 3)
    return stochastic_choice(v)